package org.lanettiesso.loader;

import org.apache.commons.io.IOUtils;

import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Set;
import java.util.jar.JarEntry;
import java.util.jar.JarFile;

/**
 * @author lan
 */
public abstract class IdentityClassLoader extends ClassLoader {

    public IdentityClassLoader(ClassLoader parent) {
        super(parent);
    }

    public IdentityClassLoader() {
    }

    private final Set<String> classpathSet = new HashSet<>();

    protected void addClasspath(String path) {
        classpathSet.add(path);
    }

    @Override
    protected Class<?> findClass(String name) throws ClassNotFoundException {
        Iterator<String> iterator = classpathSet.iterator();
        String clazzPath = name.replaceAll("\\.", "/") + ".class";
        while (iterator.hasNext()) {
            String next = iterator.next();
            File clazzFile = new File(next + "/" + clazzPath);
            if (clazzFile.exists()) {
                try {
                    byte[] bytes = IOUtils.toByteArray(new FileInputStream(clazzFile));
                    return defineClass(name, bytes, 0, bytes.length);
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        }
        try {
            iterator = classpathSet.iterator();
            while (iterator.hasNext()) {
                String path = iterator.next();
                Class<?> c = tryLoadFromJar(new File(path), name);
                if (c != null) {
                    return c;
                }
            }
        } catch (IOException ignored) {
        }
        throw new ClassNotFoundException(name);
    }

    private Class<?> tryLoadFromJar(File file, String clazzName) throws IOException {
        if (file.isDirectory()) {
            File[] files = file.listFiles();
            for (File f : files) {
                Class<?> c = tryLoadFromJar(f, clazzName);
                if (c != null) {
                    return c;
                }
            }
        } else {
            if (file.getName().endsWith(".jar")) {
                String clazzPath = clazzName.replaceAll("\\.", "/") + ".class";
                JarFile jarFile = new JarFile(file);
                JarEntry jarEntry = jarFile.getJarEntry(clazzPath);
                if (jarEntry != null) {
                    InputStream inputStream = jarFile.getInputStream(jarEntry);
                    byte[] bytes = IOUtils.toByteArray(inputStream);
                    return defineClass(clazzName, bytes, 0, bytes.length);
                }
            }
        }
        return null;
    }
}
